package org.sky.cloud.common.executor;

import cn.hutool.core.collection.ConcurrentHashSet;
import org.sky.cloud.common.utils.ThreadUtils;

import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicBoolean;

public class ThreadPoolManager {

    private Map<String, Set<ExecutorService>> executorMap;

    private Map<String, Object> locks = new ConcurrentHashMap<>(8);

    private static final ThreadPoolManager INSTANCE = new ThreadPoolManager();

    private  AtomicBoolean closed = new AtomicBoolean(false);


    static {
        INSTANCE.init();
        ThreadUtils.addShutdownHook(new Thread(ThreadPoolManager::shutdown));
    }

    public void init(){
        executorMap = new ConcurrentHashMap<>();
    }

    public static ThreadPoolManager getInstance(){
        return INSTANCE;
    }

    public void registerExecutor(String name, ExecutorService executor){
        if(closed.get()){
            throw new IllegalStateException("ThreadPoolManager is closed");
        }
        synchronized (locks.computeIfAbsent(name, k -> new Object())) {
            if(executorMap == null){
                executorMap = new HashMap<>();
            }
            Set<ExecutorService> executors = executorMap.get(name);
            if(executors == null){
                executors = new ConcurrentHashSet<>();
                executors.add(executor);
            }else{
                executors.add(executor);
            }
            executorMap.put(name, executors);
        }
    }


    public void destroy(String name){
        final Object monitor = locks.get(name);
        if (monitor == null) {
            return;
        }
        synchronized (monitor) {
            Set<ExecutorService> subResource = executorMap.get(name);
            if (subResource == null) {
                return;
            }
            for (ExecutorService executor : subResource) {
                ThreadUtils.shutdownThreadPool(executor);
            }
            executorMap.get(name).clear();
            executorMap.remove(name);
        }
    }

    public static void shutdown(){
        if(INSTANCE.closed.compareAndSet(false, true)){
            if(INSTANCE.executorMap != null){
                Set<String> keySet = INSTANCE.executorMap.keySet();
                for (String key : keySet) {
                    INSTANCE.destroy(key);
                }
            }
        }
    }
}
